//      gene_em_impl.h
//      
//      Copyright 2012 tqlong <tqlong@espada>
//      
//      This program is free software; you can redistribute it and/or modify
//      it under the terms of the GNU General Public License as published by
//      the Free Software Foundation; either version 2 of the License, or
//      (at your option) any later version.
//      
//      This program is distributed in the hope that it will be useful,
//      but WITHOUT ANY WARRANTY; without even the implied warranty of
//      MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
//      GNU General Public License for more details.
//      
//      You should have received a copy of the GNU General Public License
//      along with this program; if not, write to the Free Software
//      Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
//      MA 02110-1301, USA.

#ifndef __GENE_EM_IMPL_H
#define __GENE_EM_IMPL_H

template<typename H>
void ProbMap<H>::normalize() {
  double s = sum();
  BOOST_FOREACH(value_type& i, *this) {
    i.second = LogNumber::div(i.second, s);
  }
}

template<typename H>
double ProbMap<H>::sum() const {
  double s = LogNumber::zero();
  BOOST_FOREACH(const value_type& i, *this) {
    s = LogNumber::add(i.second, s);
  }    
  return s;
}

template<typename H>
void ProbMap<H>::print() const {
  BOOST_FOREACH(const value_type& i, *this) {
    i.first.print();
    std::cerr << " --> " << i.second << std::endl;
  }
}

template <typename M, typename D>
  double gene_em(int max_iter, M& model, const D& data) 
{
  typedef typename M::datum_t datum_t;
  typedef typename M::prob_t prob_t;

  double likelihood = -DOUBLE_INFINITY;
  for (int iter = 0; iter < max_iter; iter++) {
    MSG("-------------  iter = " << iter << "  -------------");
    //model.print();

    model.begin_accumulate();
    likelihood = LogNumber::one();
    BOOST_FOREACH(const datum_t& x, data) {
      prob_t p = model.compute_full_probability(x); // get joint probability p(x, hidden)
      likelihood = LogNumber::mul(likelihood, p.sum());
      //MSG("  prob(" << x << ")="); p.print();
      p.normalize(); // get posterior probability p(hidden|x)
      //MSG("  posterior(" << x << ")="); p.print();
      model.accumulate(x, p);
    }
    model.end_accumulate();
    MSG("  likelihood = " << likelihood);
  }
  std::cerr << "final " << model.toString();
  return likelihood;
}

template <typename M, typename GD>
  double gene_em_max(int max_iter, M& model, const GD& data) 
{
  typedef typename M::datum_t datum_t;
  typedef typename M::prob_t prob_t;
  typedef typename GD::value_type data_t;

  double likelihood = -DOUBLE_INFINITY;
  for (int iter = 0; iter < max_iter; iter++) {
    MSG("-------------  iter = " << iter << "  -------------");
    //model.print();

    model.begin_accumulate();
    likelihood = LogNumber::one();

    BOOST_FOREACH(const data_t& group, data) {
      const datum_t* max_x = NULL;
      double max_p = -DOUBLE_INFINITY;
      prob_t max_prob;
      BOOST_FOREACH(const datum_t& x, group) {
        prob_t p = model.compute_full_probability(x);
        double p_x = p.sum();
        if (p_x > max_p) {
          max_x = &x;
          max_p = p_x;
          max_prob = p;
        }
      }
      likelihood = LogNumber::mul(likelihood, max_p);
      max_prob.normalize();
      model.accumulate(*max_x, max_prob);
    }

    model.end_accumulate();
    MSG("  likelihood = " << likelihood);
  }
  std::cerr << "final " << model.toString();
  return likelihood;
}

#endif
